/* Author: Kuljit Virk

This file is self-contained interface for the input/output of the data
structures, mainly the main inputs to the code and storage of Green functions, 
over the two-time plane. The basic
strategy is to store the Green functions in the reverse chronological order
of their FIRST time argument. Each such green function is represented
by a separate array of pointers, where each pointer points to an array 
for the second time argument + the kpoint. The storage strategy allows for
a sequenced retreival of the G(t,t') functions as they are needed in the
self energy calculations.

Inputs to the code: 

time 
Efield
Kpoints
Band Hamiltonian
Interaction Hamiltonian

Best solution might be to write hdf5 file containing all this and it 
will be a single package for different inputs.

*/

#include <io.h>
#include <hdf5_hl.h>



template <class dtype>
// TODO: Explain why buff is a pointer to pointer
int read_dset(const hid_t file,const std::string dset, dtype **buff, int *idims )
{
  using namespace std;
  using namespace H5;
  string errprefix="\n\nERROR read_dset(): ";
  hsize_t dims[3];
  H5T_class_t class_id;
  size_t type_size;
  if (! H5LTfind_dataset (file,dset.c_str())) 
    {
      cerr << errprefix << " dset "<<dset<<" not found\n\n";
      exit(0);
    }
  int ndims;
  H5LTget_dataset_ndims(file,dset.c_str(),&ndims);
  H5LTget_dataset_info(file,dset.c_str(),dims, &class_id, &type_size);
  int totdims=1;
  for (int i=0; i<ndims; i++)
    {
      idims[i]   = dims[i];
      totdims   *= dims[i];
    }
  *buff = new dtype[totdims];
  H5LTread_dataset_double(file,dset.c_str(),*buff);
  return 1;
}
/*
void io::write_output_hdf5_file(const char* filename,
				const complex_matrix_function& rho)
{
  using namespace std;
  using namespace arma;
  using namespace H5;

  // Open the input HDF5 file
  hid_t file=H5Fopen(filename,H5F_ACC_RDWR,H5P_DEFAULT);
  int idims[3];
  string dset;

  // Read the time array: dims[0]x1 = array
  double *t;
  // TODO: Explain why we send pointer to the pointer t
  H5LTmake_dataset_double(file,dset,2
  int success=write_dset(file,"rho_real",rho,idims);
  if (! success ) return;
  int nt=idims[0];
  
  
}
*/
void io::read_input_hdf5_file(const char* filename,
			      arma::vec& time,
			      arma::vec& Efield,
			      arma::mat& kpoints,
			      arma::mat& EnergyBands,
			      complex_matrix_function& InteractionHamiltonian,
			      complex_matrix_function& rho0)
{
  using namespace std;
  using namespace arma;
  using namespace H5;

  // Open the input HDF5 file
  hid_t file=H5Fopen(filename,H5F_ACC_RDONLY,H5P_DEFAULT);
  
  int idims[3];
  string dset;

  // Read the time array: dims[0]x1 = array
  double *t;
  // TODO: Explain why we send pointer to the pointer t
  int success=read_dset(file,"time",&t,idims);
  if (! success ) return;
  int nt=idims[0];

  // Read the Efield array
  double *Et;
  success=read_dset(file,"Efield",&Et,idims);
  if (! success ) return;
  // Read the kpoints array
  double *kp;
  success=read_dset(file,"kpoints",&kp,idims);
  if (! success ) return;
  int nk=idims[0];

  // Read Energy Bands
  double *EB;
  success=read_dset(file,"EnergyBands",&EB,idims);
  if (! success ) return;
  
  int nb = idims[1];
  if (idims[0] != nk)
    {
      cerr << "\n\nERROR read_input_hdf5_file(): EnergyBands and kpoints arrays do not match\n\n";
      exit(0);
    }

  double *IHr,*IHi, *RHr,*RHi;
  success=read_dset(file,"InteractionHamiltonian_real",&IHr,idims);
  if (! success ) return;
  success=read_dset(file,"InteractionHamiltonian_imag",&IHi,idims);
  if (! success ) return;

  success=read_dset(file,"rho0_real",&RHr,idims);
  if (! success ) return;
  success=read_dset(file,"rho0_imag",&RHi,idims);
  if (! success ) return;


  //printf("%5d %5d %5d; %5d %5d %5d\n",idims[0],idims[1],idims[2],nt,nk,nb);
  
  // COPY DATA
  time.zeros(nt);
  Efield.zeros(nt);  
  kpoints.zeros(nk,3);
  EnergyBands.zeros(nk,nb);
  cx_mat H(nb,nb);
  H.zeros();
  InteractionHamiltonian.resize(nk,H);
  rho0.resize(nk,H);
  

  for (int it=0; it < nt ; it++)
    {
      time(it) = t[it];
      Efield(it) = Et[it];
    }
  
  for (int ik=0; ik < nk ; ik++)
    {
      // Assuming kvectors are stored as k(:,nk)
      kpoints(ik,0) = kp[ik*2];
      kpoints(ik,1) = kp[ik*2 + 1];
      for (int ib1=0; ib1 < nb; ib1++)
	{
	  // Assuming EnergyBands stored as (0..nb-1,0..nk-1)
	  EnergyBands(ik,ib1) = EB[ib1 + ik*nb];
	  for (int ib2=0; ib2 < nb; ib2++)
	    {	
	      InteractionHamiltonian[ik](ib1,ib2) = 
		complex<double>(IHr[ib2 + ib1*nb + ik*nb*nb],IHi[ib2 + ib1*nb + ik*nb*nb]);
	      rho0[ik](ib1,ib2) = 
		complex<double>(RHr[ib2 + ib1*nb + ik*nb*nb],RHi[ib2 + ib1*nb + ik*nb*nb]);
	      //printf(": %5d %5d %5d  %14.4f\n",ib2,ib1,ik,IHr[ib2 + ib1*nb + ik*nb*nb]);
	    }
	}
    }

  return;
}

  
// This function reads kpoints stored in a text file: it will be deprecated
void io::read_kpoints(const char* filename, arma::mat& kpoints, double& dkx, int& nx)
{
  using namespace std;
  ifstream kfile;
  vector<double> kx,ky;
  double px,py;
  kfile.open(filename);
  if (kfile.is_open())
    {
      while ( kfile.good() ) 
	{
	  // Read the stream
	  kfile >> px>>py;
	  // When EOF is reached, the eof flag is set 
	  // px,py are not updated. So we must check that the field.good()
	  // is true otherwise we push the last value to vector twice
	  if (kfile.good())
	    {
	      kx.push_back(px);
	      ky.push_back(py);
	    }
	}
    }
  else
    std::cout<<"ERROR: io::read_kpoints()  file "<<filename<<" did not open\n ";
  kfile.close();
  
  int nk=kx.size();
  kpoints.zeros(nk,3);
  for(int i=0;i<nk;i++)
    {
      kpoints(i,0) = kx[i];
      kpoints(i,1) = ky[i];
    }
  
  // HARDWIRED: kx changes the fastest in the file
  dkx = kx[1]-kx[0];

  while (std::fabs(ky[nx]-ky[0])<1.0e-8)
    nx++;
}

// This function reads (time,Field) pairs from a text file: it will be deprecated
void io::read_fields(const char* filename, arma::vec& time, arma::vec& Efield)
{
  using namespace std;
  ifstream field;
  vector<double> t,Ef;
  double t0,E;
  field.open(filename);
  if (field.is_open())
    {
      while ( field.good() ) 
	{
	  // Read the stream
	  field >> t0>>E;
	  // When EOF is reached, the eof flag is set 
	  // t0 and E are not updated. So we must check that the field.good()
	  // is true otherwise we push the last value to vector twice
	  if (field.good())
	    {
	      t.push_back(t0);
	      Ef.push_back(E);
	    }
	}
    }
  else
    std::cout<<"ERROR: io::read_fields()  file "<<filename<<" did not open\n ";
  field.close();

  int nt=t.size();
  time.zeros(nt);
  Efield.zeros(nt);
  for (int i=0; i<nt; i++)
    {
      time(i)   = t[i];
      Efield(i) = Ef[i]; 
    }
}

void io::set_filenames(h5files_container& files)
{
	using namespace H5;
	files=h5files_container(4);

	std::cout<<"\nset_filenames():\n";  	

	files[0].name = "g_gret.h5";
	files[1].name = "g_less.h5";
	files[2].name = "self_energy_gret.h5";
	files[3].name = "self_energy_less.h5";
}

void io::disk_storage_initialization(
	h5files_container& files,
	const struct parameters& params)
{
	using namespace std;
	using namespace H5;

	std::cout<<"\ndisk_storage_initialization():\n";


	int one_chunk_bit= params.ntmem * params.num_bands * params.num_bands;
	int N = params.num_mesh_points * one_chunk_bit;
	int error;

	for ( int i=0; i < files.size(); i++)
	{

		files[i].file_id = H5Fcreate(files[i].name.c_str(), 
			H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
		/* to store N complex numbers as:  [1....N][1....N]
		*                                   real    imag   */
			files[i].chunk_dims[0] = 2*one_chunk_bit;
		files[i].chunk_dims[1] = 1;
		files[i].dset_dims[0]  = 2*N;
		files[i].dset_dims[1]  = params.ntmem;

		/* Property lists common to all files */
		/* Create the dataspace and the chunked dataset */
		files[i].space_id = H5Screate_simple(2, files[i].dset_dims, NULL);
		/* Create a dataset creation property list and set it to use chunking
		* with a chunk size of chunk_dims */
			files[i].dcpl_id = H5Pcreate(H5P_DATASET_CREATE);
		H5Pset_chunk(files[i].dcpl_id, 2, files[i].chunk_dims);
		/* Create the dataset */
		files[i].dset_id = H5Dcreate(files[i].file_id,"values",
			H5T_NATIVE_DOUBLE,
			files[i].space_id,H5P_DEFAULT,
			files[i].dcpl_id,
			H5P_DEFAULT);

		printf("file:  %10s:  chunk_dims = %3d  %3d\n",files[i].name,
			files[i].chunk_dims[0],files[i].chunk_dims[1]);

	}

}

void io::dump(const h5_chunked_file_info& file,
	const complex_matrix_function& M, 
	const struct parameters& params,
	const int t1_index)
{
	using namespace std;
	using namespace H5;
	int n_elem = M.size();
	int N      = params.num_mesh_points*params.num_bands*params.num_bands;
	double *R  = new double[2*N];

	int itau, ik, i,ib1,ib2,j;
	hsize_t start[2]    = {t1_index*2*N,0};
	hsize_t mem_dims[1] = {2*N};
	hsize_t count[2]    = {mem_dims[0],1};

	hid_t mspace_id;
	hid_t space_id;
	htri_t er;
	herr_t err;

/* Memory space id */
	mspace_id = H5Screate_simple(1, mem_dims, NULL);
	i=0;

/*

Time grid:
----------.
--------. . tau5
------. . . tau4
----. . . . tau3
--. . . . . tau2
. . . . . . tau1
t1  t3  t5

Array in memory has tau changing faster than k
+------------+------------+.....+---------------+-----------+....
| k=0, tau=1 | k=0, tau=2 |.....| k=0, tau=ntau | k=1 tau=0 |....
+------------+------------+.....+---------------+-----------+....

Chunks in the file are stored in columns shown below where k changes
faster than the times t1,t3 along the rows below. The columns, tau1, tau2,...
represent each of the horizontal lines. They are different chunks
in the file.

tau=1       tau=2              
+------------+------------+.....+---------------+-----------+....
t1    |            |
+------------+------------+.....+---------------+-----------+....
t2    |            | 
+------------+------------+.....+---------------+-----------+....
t3    |            | 
+------------+------------+.....+---------------+-----------+....
t4    |            | 
+------------+------------+.....+---------------+-----------+....

An array of N elements is stored as [Real part | Imaginary part]

*/


/* Go along the vertical direction and store each row tau1, tau2,...  */
for (itau=0; itau < params.ntmem; itau++)
{
	// TODO: handling boundary? 
	j=0;
	for ( ik=0; ik < params.num_mesh_points; ik++ )
	{
		i = ik * params.ntmem + itau;
		for ( ib1 = 0; ib1 < params.num_bands; ib1++)
			for ( ib2 = 0; ib2 < params.num_bands; ib2++)
		{
			R[j]     = real(M[i](ib1,ib2));
			R[j + N] = imag(M[i](ib1,ib2));
			j++;
		}
	}
	// Now perform a single write opearation with R.
	/* The column is itau */
	start[1] = itau;
	space_id  = file.space_id; //try this if trouble: H5Dget_space(file.dset_id);
	H5Sselect_hyperslab(space_id, H5S_SELECT_SET, start, NULL, count, NULL);
	/* Write to the dataset */
	H5Dwrite(file.dset_id,H5T_NATIVE_DOUBLE,mspace_id,space_id,H5P_DEFAULT,R);
}
}

void io::load(const h5_chunked_file_info& file,
	const parameters& params,
	const int itau,
	complex_matrix_function& M)
{
	using namespace std;
	using namespace H5;
// read chunk itau and fill a vector 
// the most recent chunks in itau line are from 0...itau-1

	int nbit = params.num_mesh_points * params.num_bands * params.num_bands;
	int N    = (itau+1)*nbit;

/* Create the memory dataspace */
	hid_t mspace_id;
	hsize_t mem_dims[1] = {2*N};
	mspace_id = H5Screate_simple(1, mem_dims, NULL);

	hsize_t start[2];
	hsize_t count[2] ={2*N,1};
	double *buffer = new double[2*N];

/* Real part */
	start[0] = 0;
	start[1] = itau;

	H5Sselect_hyperslab(file.space_id, H5S_SELECT_SET, start, NULL, count, NULL);
	H5Dread(file.dset_id, H5T_NATIVE_DOUBLE,
		mspace_id, file.space_id, H5P_DEFAULT, buffer);

	int i,j,jj, ib1,ib2,ik;
	int nbsq = params.num_bands*params.num_bands;

	for ( i = 0; i < itau ; i++ )
	{
		jj = i*params.num_mesh_points;
		for (ik=0; ik<params.num_mesh_points;ik++)
		{
			j = ik*params.ntmem + i;
			for ( ib1 = 0; ib1 < params.num_bands ; ib1++ )
				for ( ib2 = 0; ib2 < params.num_bands ; ib2++ )
			{
				M[j](ib1,ib2) = dcomplex(buffer[jj],buffer[jj+nbit]);
				jj++;
			}
		}
	}

	jj = 0;


}
